from PIL import Image, ImageDraw, ImageFont

def concatenate_images_with_labels(image_path1, image_path2, output_path, direction='horizontal'):
    # 打开图片
    img1 = Image.open(image_path1)
    img2 = Image.open(image_path2)

    # 确定拼接方向
    if direction == 'horizontal':
        new_width = img1.width + img2.width
        new_height = max(img1.height, img2.height)
        new_image = Image.new('RGBA', (new_width, new_height), (255, 255, 255, 255))
        new_image.paste(img1, (0, 0))
        new_image.paste(img2, (img1.width, 0))
        # 标记位置
        label_positions = [(10, 10), (img1.width + 10, 10)]
    elif direction == 'vertical':
        new_width = max(img1.width, img2.width)
        new_height = img1.height + img2.height
        new_image = Image.new('RGBA', (new_width, new_height), (255, 255, 255, 255))
        new_image.paste(img1, (0, 0))
        new_image.paste(img2, (0, img1.height))
        # 标记位置
        label_positions = [(10, 10), (10, img1.height + 10)]
    else:
        raise ValueError("方向必须是 'horizontal' 或 'vertical'")

    # 添加标记
    draw = ImageDraw.Draw(new_image)

    try:
        # 使用 Pillow 自带字体
        font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
        font = ImageFont.truetype(font_path, 30)
    except IOError:
        raise IOError("无法加载字体，请确认路径是否正确。")


    draw.text(label_positions[0], "a", fill="black", font=font)
    draw.text(label_positions[1], "b", fill="black", font=font)

    # 保存拼接结果
    new_image.save(output_path)
    print(f"图片已保存至 {output_path}")

# 示例用法
concatenate_images_with_labels('Trade_off/a_frozen_layers.png', 'Trade_off/b_train_epochs.png', 'output_with_labels.png', direction='vertical')
